#!/bin/bash

mstack_script_dir=$(dirname "$(realpath "${BASH_SOURCE[0]}")")
source $mstack_script_dir/common_variables.sh


read -s -p "请输入节点 root 密码: " common_passwd
echo  # 打印换行符

## 安装sshpass
function sshpass_install() {
    # 检查是否安装
    if ! command -v sshpass &> /dev/null; then
        echo -e "${RED}sshpass 未安装，正在尝试安装...${NC}"
        sudo dpkg -i  $mstack_packet_archives_dir/sshpass_1.09-1_amd64.deb
    else
        echo -e "${GREEN}sshpass 已安装。${NC}"
    fi
}

## 创建密钥对并添加到 authorized_keys 文件
function create_secret_key() {
    # 第二步：检查 SSH 密钥
    SSH_DIR="/root/.ssh"
    if [ ! -d "$SSH_DIR" ]; then
      echo -e "${YELLOW}创建 .ssh 目录。${NC}"
      mkdir -p $SSH_DIR
      chmod 700 $SSH_DIR
    fi

    if [ ! -f "$SSH_DIR/id_rsa" ] || [ ! -f "$SSH_DIR/id_rsa.pub" ]; then
      echo -e "${YELLOW}创建 SSH 密钥对。${NC}"
      ssh-keygen -t rsa -b 4096 -f "$SSH_DIR/id_rsa" -q -N ""
    fi

    PUB_KEY=$(cat "$SSH_DIR/id_rsa.pub")
    AUTHORIZED_KEYS="$SSH_DIR/authorized_keys"
    if ! grep -q "$PUB_KEY" "$AUTHORIZED_KEYS"; then
      echo -e "${YELLOW}添加公钥到 authorized_keys 文件。${NC}"
      echo "$PUB_KEY" >> $AUTHORIZED_KEYS
      chmod 600 $AUTHORIZED_KEYS
    fi

    # 获取执行脚本主机的 id_rsa.pub 文件路径
    public_key_file="$HOME/.ssh/id_rsa.pub"
    echo -e "${BLUE}当前主机公钥地址 $HOME/.ssh/id_rsa.pub${NC}"

    # 检查 id_rsa.pub 文件是否存在
    if [[ ! -f $public_key_file ]]; then
        echo -e "${RED}错误: 公钥文件 $public_key_file 不存在.${NC}"
        exit 1
    fi

    # 读取公钥内容
    public_key=$(cat "$public_key_file")
}



function slave_add_master_public_key() {
    sshpass_install
    create_secret_key

    # 读取 hosts 文件并处理
    for ip in $(awk '{print $1}' $mstack_tools_mstack_dir/mstack_all_nodes.txt | sed -r '/^[ \t]*($|#)/d'); do
        # 通过 grep 获取对应的主机名
        name=$(grep "$ip" $mstack_tools_mstack_dir/mstack_all_nodes.txt| awk '{print $2}' | sed -r '/^[ \t]*($|#)/d')

            # 输出结果
            echo -e "${CYAN}mstack节点 IP: $ip, 节点名称 :$name${NC}"

            echo -e "${CYAN}尝试连接到 $ip...${NC}"

            # 使用 sshpass 尝试连接
            sshpass -p "$common_passwd" ssh -o StrictHostKeyChecking=no root@$ip "hostname; uptime"

            # 检查连接是否成功
            if [ $? -ne 0 ]; then
                echo -e "${RED}连接到 $ip 失败，请输入新的密码:${NC}"
                read -s -p "请输入 $ip 的 root 用户密码: " new_passwd
                echo  # 打印换行符
                # 尝试再次连接
                sshpass -p "$new_passwd" ssh -o StrictHostKeyChecking=no root@$ip "hostname; uptime"

                # 检查新的密码连接是否成功
                if [ $? -ne 0 ]; then
                    echo -e "${RED}连接到 $ip 失败，可能密码错误。${NC}"
                    continue
                else
                    # 检查 authorized_keys 文件
                    sshpass -p "$new_passwd" ssh -o StrictHostKeyChecking=no root@$ip "grep -qF '$public_key' /root/.ssh/authorized_keys"

                    # 检查命令执行结果
                    if [ $? -ne 0 ]; then
                        echo -e "${YELLOW}公钥未找到，将添加到 /root/.ssh/authorized_keys 文件。${NC}"
                        # 如果文件不存在，则创建
                        sshpass -p "$new_passwd" ssh -o StrictHostKeyChecking=no root@$ip "mkdir -p /root/.ssh; echo '$public_key' >> /root/.ssh/authorized_keys; chmod 600 /root/.ssh/authorized_keys"
                        echo -e "${GREEN}公钥已添加。${NC}"
                    else
                        echo -e "${GREEN}公钥已存在于 /root/.ssh/authorized_keys 文件中。${NC}"
                    fi
                fi
            else
                # 已连接成功，检查公钥是否已存在
                sshpass -p "$common_passwd" ssh -o StrictHostKeyChecking=no root@$ip "grep -qF '$public_key' /root/.ssh/authorized_keys"
                if [ $? -ne 0 ]; then
                    echo -e "${YELLOW}公钥未找到，将添加到 /root/.ssh/authorized_keys 文件。${NC}"
                    # 如果文件不存在，则创建
                    sshpass -p "$common_passwd" ssh -o StrictHostKeyChecking=no root@$ip "mkdir -p /root/.ssh; echo '$public_key' >> /root/.ssh/authorized_keys; chmod 600 /root/.ssh/authorized_keys"
                    echo -e "${GREEN}公钥已添加。${NC}"
                else
                    echo -e "${GREEN}公钥已存在于 /root/.ssh/authorized_keys 文件中。${NC}"
                fi
            fi
    done
}

slave_add_master_public_key

